import os
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report
from sklearn.metrics import roc_curve, auc
import numpy as np
import joblib
import matplotlib.pyplot as plt

def load_data(data_dir, n_frames=2000):
    asd_data = []
    td_data = []

    for filename in os.listdir(data_dir['ASD']):
        if filename.endswith('.csv'):
            df = pd.read_csv(os.path.join(data_dir['ASD'], filename))
            # 补零到n_frames
            if len(df) < n_frames:
                pad_len = n_frames - len(df)
                pad = pd.DataFrame(np.zeros((pad_len, df.shape[1])), columns=df.columns)
                df = pd.concat([df, pad], ignore_index=True)
            asd_data.append(df.iloc[:n_frames])

    for filename in os.listdir(data_dir['TD']):
        if filename.endswith('.csv'):
            df = pd.read_csv(os.path.join(data_dir['TD'], filename))
            if len(df) < n_frames:
                pad_len = n_frames - len(df)
                pad = pd.DataFrame(np.zeros((pad_len, df.shape[1])), columns=df.columns)
                df = pd.concat([df, pad], ignore_index=True)
            td_data.append(df.iloc[:n_frames])

    return asd_data, td_data

def preprocess_data(asd_data, td_data):
    asd_labels = [1] * len(asd_data)
    td_labels = [0] * len(td_data)

    all_data = asd_data + td_data
    all_labels = asd_labels + td_labels

    X = []
    for data in all_data:
        # 只取Gaze_X, Gaze_Y, Expression并展平
        features = data[['Gaze_X', 'Gaze_Y', 'Expression']].values.flatten()
        X.append(features)
    X = np.array(X)
    y = np.array(all_labels)

    return train_test_split(X, y, test_size=0.2, stratify=y, random_state=42)

def train_model(X_train, y_train):
    model = RandomForestClassifier(n_estimators=100, random_state=42)
    model.fit(X_train, y_train)
    return model

def evaluate_model(model, X_test, y_test):
    y_pred = model.predict(X_test)
    report = classification_report(y_test, y_pred)
    print(report)

def plot_roc(model, X_test, y_test):
    # 获取预测概率
    if hasattr(model, "predict_proba"):
        y_score = model.predict_proba(X_test)[:, 1]
    else:
        y_score = model.decision_function(X_test)
    fpr, tpr, thresholds = roc_curve(y_test, y_score)
    roc_auc = auc(fpr, tpr)

    plt.figure()
    plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (AUC = {roc_auc:.2f})')
    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver Operating Characteristic')
    plt.legend(loc="lower right")
    plt.show()

if __name__ == "__main__":
    data_dir = {
        'ASD': r'D:\毕设\ASD',
        'TD': r'D:\毕设\TD'
    }

    asd_data, td_data = load_data(data_dir)
    X_train, X_test, y_train, y_test = preprocess_data(asd_data, td_data)
    model = train_model(X_train, y_train)
    evaluate_model(model, X_test, y_test)

    # 自动创建model文件夹
    model_dir = r'd:\毕设\asd-classification-project\model'
    os.makedirs(model_dir, exist_ok=True)

    # 保存模型
    joblib.dump(model, os.path.join(model_dir, 'trained_model.pkl'))

    # 保存测试集
    test_df = pd.DataFrame(X_test)
    test_df['label'] = y_test
    test_df.to_csv(os.path.join(model_dir, 'test_data.csv'), index=False)

    # 绘制ROC曲线
    plot_roc(model, X_test, y_test)